import torch
import random
from mesh_coverage_loss import _p2e_logproba

seed = 66
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

n_f = 10000
n_s = 20
n_e = 70000

sxy = torch.randn(n_f, n_s, 2, dtype=torch.float32).cuda().requires_grad_(True)
oxy = torch.randn(n_e, 2, dtype=torch.float32).cuda().detach().requires_grad_(True)
theta = torch.rand(n_e) * torch.pi

invcov_root = torch.stack([torch.cos(theta), -torch.sin(theta), torch.sin(theta), torch.cos(theta)], dim=1).reshape(n_e, 2, 2).contiguous().cuda().detach().requires_grad_(True)
invcov = torch.bmm(invcov_root, invcov_root.permute(0, 2, 1))
invcov.retain_grad()
logdet_invcov_root = torch.logdet(invcov_root)
logdet_invcov_root.retain_grad()
fids = torch.randint(0, n_f, (n_e, ), dtype=torch.int32).cuda()

print(fids)
logproba_a = _p2e_logproba.apply(sxy, oxy, invcov, logdet_invcov_root, fids)
print(logproba_a.shape)
loss = logproba_a.mean()
loss.backward()
print(sxy.grad)
print(oxy.grad)
print(invcov.grad)
print(logdet_invcov_root.grad)